# -*- coding: utf-8 -*-
"""
知识库管理系统 - 连接pgvector数据库实现增删改查
"""

import psycopg2
import requests
import json
import numpy as np
from typing import List, Dict, Optional, Tuple
from config import DATABASE_CONFIG, OLLAMA_CONFIG, VECTOR_DIMENSION, TABLE_NAME


class KnowledgeBase:
    """知识库管理类"""
    
    def __init__(self, table_name=None):
        """
        初始化知识库连接
        
        Args:
            table_name: 指定表名，如果不指定则使用配置文件中的默认表名
        """
        self.conn = None
        self.cursor = None
        self.table_name = table_name or TABLE_NAME
        self.connect_database()
        self.check_and_create_table()
    
    def connect_database(self):
        """连接PostgreSQL数据库"""
        try:
            self.conn = psycopg2.connect(**DATABASE_CONFIG)
            self.cursor = self.conn.cursor()
            print("✅ 数据库连接成功")
        except Exception as e:
            print(f"❌ 数据库连接失败: {e}")
            raise
    
    def check_and_create_table(self):
        """检查表是否存在，不存在则创建"""
        try:
            # 检查表是否存在
            if self.table_exists():
                print(f"✅ 表 '{self.table_name}' 已存在，跳过创建")
                return
            
            print(f"📋 表 '{self.table_name}' 不存在，开始创建...")
            self.create_table()
            
        except Exception as e:
            print(f"❌ 检查/创建表失败: {e}")
            raise
    
    def table_exists(self):
        """检查表是否存在"""
        try:
            check_sql = """
            SELECT EXISTS (
                SELECT FROM information_schema.tables 
                WHERE table_schema = 'public' 
                AND table_name = %s
            );
            """
            self.cursor.execute(check_sql, (self.table_name,))
            return self.cursor.fetchone()[0]
        except Exception as e:
            print(f"❌ 检查表存在性失败: {e}")
            return False
    
    def create_table(self):
        """创建知识库表"""
        try:
            # 启用必要的扩展
            self.cursor.execute("CREATE EXTENSION IF NOT EXISTS vector;")
            self.cursor.execute("CREATE EXTENSION IF NOT EXISTS \"uuid-ossp\";")
            
            # 创建表
            create_table_sql = f"""
            CREATE TABLE "{self.table_name}" (
                "id" uuid NOT NULL DEFAULT uuid_generate_v4() PRIMARY KEY,
                "text" text COLLATE "pg_catalog"."default" NOT NULL,
                "metadata" jsonb DEFAULT '{{}}',
                "embedding" vector({VECTOR_DIMENSION})
            );
            """
            self.cursor.execute(create_table_sql)
            
            # 创建向量索引以提高查询性能
            index_sql = f"""
            CREATE INDEX "{self.table_name}_embedding_idx" 
            ON "{self.table_name}" USING ivfflat ("embedding" vector_cosine_ops) 
            WITH (lists = 100);
            """
            self.cursor.execute(index_sql)
            
            self.conn.commit()
            print(f"✅ 表 '{self.table_name}' 创建完成")
            
        except Exception as e:
            print(f"❌ 创建表失败: {e}")
            self.conn.rollback()
            raise
    
    def get_embedding(self, text: str) -> List[float]:
        """获取文本嵌入向量"""
        try:
            url = f"{OLLAMA_CONFIG['base_url']}/api/embeddings"
            payload = {
                "model": OLLAMA_CONFIG['embedding_model'],
                "prompt": text
            }
            
            response = requests.post(url, json=payload, timeout=30)
            response.raise_for_status()
            
            result = response.json()
            embedding = result.get('embedding', [])
            
            if not embedding:
                raise ValueError("未获取到有效的嵌入向量")
            
            return embedding
        except Exception as e:
            print(f"❌ 获取嵌入向量失败: {e}")
            raise
    
    def format_vector(self, embedding: List[float]) -> str:
        """将Python列表格式化为pgvector字符串格式"""
        return '[' + ','.join(map(str, embedding)) + ']'
    
    def add_knowledge(self, text: str, metadata: Dict = None) -> str:
        """
        添加知识到数据库
        
        Args:
            text: 知识文本内容
            metadata: 元数据（可选）
        
        Returns:
            str: 插入记录的UUID
        """
        try:
            # 获取嵌入向量
            embedding = self.get_embedding(text)
            
            # 准备元数据
            if metadata is None:
                metadata = {}
            
            # 插入数据 - 使用字符串格式的向量
            insert_sql = f"""
            INSERT INTO "{self.table_name}" ("text", "embedding", "metadata")
            VALUES (%s, %s::vector, %s)
            RETURNING "id";
            """
            
            vector_str = self.format_vector(embedding)
            self.cursor.execute(insert_sql, (text, vector_str, json.dumps(metadata)))
            record_id = self.cursor.fetchone()[0]
            self.conn.commit()
            
            print(f"✅ 知识添加成功，ID: {record_id}")
            return str(record_id)
        
        except Exception as e:
            print(f"❌ 添加知识失败: {e}")
            self.conn.rollback()
            raise
    
    def search_knowledge(self, query: str, limit: int = 5, similarity_threshold: float = 0.5) -> List[Dict]:
        """
        搜索相似知识
        
        Args:
            query: 查询文本
            limit: 返回结果数量限制
            similarity_threshold: 相似度阈值（0-1）
        
        Returns:
            List[Dict]: 搜索结果列表
        """
        try:
            # 获取查询向量
            query_embedding = self.get_embedding(query)
            query_vector_str = self.format_vector(query_embedding)
            
            # 执行相似度搜索 - 修复向量类型转换问题
            search_sql = f"""
            SELECT 
                "id", 
                "text", 
                "metadata",
                1 - ("embedding" <=> %s::vector) as similarity
            FROM "{self.table_name}"
            WHERE 1 - ("embedding" <=> %s::vector) > %s
            ORDER BY "embedding" <=> %s::vector
            LIMIT %s;
            """
            
            self.cursor.execute(search_sql, (
                query_vector_str, 
                query_vector_str, 
                similarity_threshold,
                query_vector_str,
                limit
            ))
            
            results = []
            for row in self.cursor.fetchall():
                results.append({
                    'id': str(row[0]),
                    'text': row[1],
                    'metadata': row[2],
                    'similarity': float(row[3])
                })
            
            print(f"✅ 搜索完成，找到 {len(results)} 条相关知识")
            return results
        
        except Exception as e:
            print(f"❌ 搜索知识失败: {e}")
            raise
    
    def get_knowledge_by_id(self, knowledge_id: str) -> Optional[Dict]:
        """
        根据ID获取知识
        
        Args:
            knowledge_id: 知识UUID
        
        Returns:
            Dict: 知识详情，如果不存在返回None
        """
        try:
            select_sql = f"""
            SELECT "id", "text", "metadata" 
            FROM "{self.table_name}"
            WHERE "id" = %s;
            """
            
            self.cursor.execute(select_sql, (knowledge_id,))
            row = self.cursor.fetchone()
            
            if row:
                return {
                    'id': str(row[0]),
                    'text': row[1],
                    'metadata': row[2]
                }
            else:
                print(f"❌ 未找到ID为 {knowledge_id} 的知识")
                return None
        
        except Exception as e:
            print(f"❌ 获取知识失败: {e}")
            raise
    
    def update_knowledge(self, knowledge_id: str, text: str = None, metadata: Dict = None) -> bool:
        """
        更新知识内容
        
        Args:
            knowledge_id: 知识UUID
            text: 新的文本内容（可选）
            metadata: 新的元数据（可选）
        
        Returns:
            bool: 更新是否成功
        """
        try:
            # 检查记录是否存在
            if not self.get_knowledge_by_id(knowledge_id):
                return False
            
            updates = []
            params = []
            
            # 如果更新内容，需要重新计算嵌入向量
            if text is not None:
                embedding = self.get_embedding(text)
                vector_str = self.format_vector(embedding)
                updates.extend(['"text" = %s', '"embedding" = %s::vector'])
                params.extend([text, vector_str])
            
            # 如果更新元数据
            if metadata is not None:
                updates.append('"metadata" = %s')
                params.append(json.dumps(metadata))
            
            if not updates:
                print("❌ 没有提供要更新的内容")
                return False
            
            # 执行更新
            update_sql = f"""
            UPDATE "{self.table_name}"
            SET {', '.join(updates)}
            WHERE "id" = %s;
            """
            params.append(knowledge_id)
            
            self.cursor.execute(update_sql, params)
            self.conn.commit()
            
            print(f"✅ 知识 ID:{knowledge_id} 更新成功")
            return True
        
        except Exception as e:
            print(f"❌ 更新知识失败: {e}")
            self.conn.rollback()
            raise
    
    def delete_knowledge(self, knowledge_id: str) -> bool:
        """
        删除知识
        
        Args:
            knowledge_id: 知识UUID
        
        Returns:
            bool: 删除是否成功
        """
        try:
            # 检查记录是否存在
            if not self.get_knowledge_by_id(knowledge_id):
                return False
            
            delete_sql = f'DELETE FROM "{self.table_name}" WHERE "id" = %s;'
            self.cursor.execute(delete_sql, (knowledge_id,))
            self.conn.commit()
            
            print(f"✅ 知识 ID:{knowledge_id} 删除成功")
            return True
        
        except Exception as e:
            print(f"❌ 删除知识失败: {e}")
            self.conn.rollback()
            raise
    
    def search_by_metadata(self, metadata_filter: Dict, limit: int = 10) -> List[Dict]:
        """
        根据元数据搜索知识
        
        Args:
            metadata_filter: 元数据过滤条件
            limit: 返回结果数量限制
        
        Returns:
            List[Dict]: 搜索结果列表
        """
        try:
            # 构建元数据查询条件
            conditions = []
            params = []
            
            for key, value in metadata_filter.items():
                conditions.append(f'"metadata" ->> %s = %s')
                params.extend([key, str(value)])
            
            where_clause = " AND ".join(conditions) if conditions else "TRUE"
            
            search_sql = f"""
            SELECT "id", "text", "metadata"
            FROM "{self.table_name}"
            WHERE {where_clause}
            ORDER BY "id"
            LIMIT %s;
            """
            params.append(limit)
            
            self.cursor.execute(search_sql, params)
            
            results = []
            for row in self.cursor.fetchall():
                results.append({
                    'id': str(row[0]),
                    'text': row[1],
                    'metadata': row[2]
                })
            
            print(f"✅ 元数据搜索完成，找到 {len(results)} 条知识")
            return results
        
        except Exception as e:
            print(f"❌ 元数据搜索失败: {e}")
            raise
    
    def list_all_knowledge(self, limit: int = 10, offset: int = 0) -> List[Dict]:
        """
        列出所有知识（分页）
        
        Args:
            limit: 每页数量
            offset: 偏移量
        
        Returns:
            List[Dict]: 知识列表
        """
        try:
            list_sql = f"""
            SELECT "id", "text", "metadata"
            FROM "{self.table_name}"
            ORDER BY "id"
            LIMIT %s OFFSET %s;
            """
            
            self.cursor.execute(list_sql, (limit, offset))
            
            results = []
            for row in self.cursor.fetchall():
                results.append({
                    'id': str(row[0]),
                    'text': row[1][:100] + '...' if len(row[1]) > 100 else row[1],  # 截断长文本
                    'metadata': row[2]
                })
            
            print(f"✅ 获取到 {len(results)} 条知识记录")
            return results
        
        except Exception as e:
            print(f"❌ 列出知识失败: {e}")
            raise
    
    def get_knowledge_count(self) -> int:
        """获取知识总数"""
        try:
            count_sql = f'SELECT COUNT(*) FROM "{TABLE_NAME}";'
            self.cursor.execute(count_sql)
            count = self.cursor.fetchone()[0]
            return count
        except Exception as e:
            print(f"❌ 获取知识总数失败: {e}")
            raise
    
    def close(self):
        """关闭数据库连接"""
        if self.cursor:
            self.cursor.close()
        if self.conn:
            self.conn.close()
        print("✅ 数据库连接已关闭")


def main():
    """演示使用示例"""
    kb = KnowledgeBase()
    
    try:
        # 1. 添加知识
        print("\n=== 添加知识示例 ===")
        id1 = kb.add_knowledge(
            "Python是一种高级编程语言，广泛用于数据科学和机器学习",
            {"category": "编程", "language": "Python"}
        )
        
        id2 = kb.add_knowledge(
            "PostgreSQL是一个功能强大的开源关系型数据库",
            {"category": "数据库", "type": "关系型"}
        )
        
        # 2. 搜索知识
        print("\n=== 搜索知识示例 ===")
        results = kb.search_knowledge("Python编程语言", limit=3)
        for result in results:
            print(f"ID: {result['id']}, 相似度: {result['similarity']:.3f}")
            print(f"内容: {result['text']}")
            print("---")
        
        # 2.5. 根据元数据搜索
        print("\n=== 元数据搜索示例 ===")
        metadata_results = kb.search_by_metadata({"category": "编程"}, limit=5)
        for result in metadata_results:
            print(f"ID: {result['id']}")
            print(f"内容: {result['text'][:50]}...")
            print(f"分类: {result['metadata'].get('category')}")
            print("---")
        
        # 3. 获取知识详情
        print("\n=== 获取知识详情示例 ===")
        knowledge = kb.get_knowledge_by_id(id1)
        if knowledge:
            print(f"内容: {knowledge['text']}")
            print(f"元数据: {knowledge['metadata']}")
        
        # 4. 更新知识
        print("\n=== 更新知识示例 ===")
        kb.update_knowledge(
            id1, 
            text="Python是一种简单易学的高级编程语言，在AI领域应用广泛",
            metadata={"category": "编程", "language": "Python", "field": "AI"}
        )
        
        # 5. 列出所有知识
        print("\n=== 列出知识示例 ===")
        all_knowledge = kb.list_all_knowledge(limit=5)
        for knowledge in all_knowledge:
            print(f"ID: {knowledge['id']}, 内容: {knowledge['text']}")
        
        # 6. 获取总数
        print(f"\n=== 知识总数: {kb.get_knowledge_count()} ===")
        
        # 7. 删除知识
        print("\n=== 删除知识示例 ===")
        kb.delete_knowledge(id2)
        
    except Exception as e:
        print(f"❌ 操作失败: {e}")
    
    finally:
        kb.close()


if __name__ == "__main__":
    main()
